import networkx as nx
from pulp import *
from collections import deque

# 管网节点类型常量
VALVE = "VALVE"
USER = "USER"
SOURCE = "SOURCE"
LEAK = "LEAK"

def build_gas_network():
    """构建示例燃气管网"""
    G = nx.DiGraph()
  
    # 添加节点（带属性）
    G.add_node("S1", type=SOURCE, pressure=1.0)
    G.add_node("V1", type=VALVE)
    G.add_node("V2", type=VALVE)
    G.add_node("V3", type=VALVE)
    G.add_node("U1", type=USER, user_weight=10)  # 医院
    G.add_node("U2", type=USER, user_weight=5)   # 学校
    G.add_node("U3", type=USER, user_weight=1)   # 居民区
    G.add_node("L1", type=LEAK)                  # 泄漏点
  
    # 添加带权重的边（管径×压力）
    G.add_edge("S1", "V1", weight=0.8*1.0)
    G.add_edge("V1", "U1", weight=0.6*0.8)
    G.add_edge("V1", "V2", weight=0.5*0.7)
    G.add_edge("V2", "U2", weight=0.4*0.6)
    G.add_edge("V2", "V3", weight=0.4*0.5)
    G.add_edge("V3", "U3", weight=0.3*0.4)
    G.add_edge("V3", "L1", weight=0.3*0.3)
  
    return G

def get_upstream_neighbors(G, node):
    """基于压力梯度获取上游邻居节点"""
    return [u for u in G.predecessors(node) 
            if G.nodes[u]["pressure"] > G.nodes[node]["pressure"]]

def get_downstream_neighbors(G, node):
    """基于压力梯度获取下游邻居节点"""
    return [v for v in G.successors(node) 
            if G.nodes[v]["pressure"] < G.nodes[node]["pressure"]]

def find_critical_valves(G, leak_nodes):
    """通过反向BFS定位关键阀门"""
    critical_valves = set()
    for leak in leak_nodes:
        queue = deque([leak])
        visited = set([leak])
        while queue:
            current = queue.popleft()
            for neighbor in get_upstream_neighbors(G, current):
                if neighbor not in visited:
                    visited.add(neighbor)
                    if G.nodes[neighbor]["type"] == VALVE:
                        critical_valves.add(neighbor)
                    else:
                        queue.append(neighbor)
    return critical_valves

def calculate_impact(G, valve):
    """计算关闭阀门的影响范围（用户权重和）"""
    impact = 0
    queue = deque([valve])
    visited = set([valve])
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("user_weight"):
            impact += G.nodes[node]["user_weight"]
        for neighbor in get_downstream_neighbors(G, node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return impact

def get_leak_paths(G, leak_node):
    """获取所有通向泄漏点的路径阀门集合"""
    paths = []
    for source in [n for n, data in G.nodes(data=True) if data["type"] == SOURCE]:
        for path in nx.all_simple_paths(G, source, leak_node):
            # 仅提取路径中的阀门节点
            path_valves = [node for node in path if G.nodes[node]["type"] == VALVE]
            if path_valves:
                paths.append(set(path_valves))
    return paths

def optimize_valve_shutdown(G, leak_nodes):
    """主优化函数"""
    # 步骤1：定位关键阀门
    critical_valves = find_critical_valves(G, leak_nodes)
    print(f"定位到关键阀门: {critical_valves}")
  
    # 步骤2：计算各阀门影响值
    impacts = {valve: calculate_impact(G, valve) for valve in critical_valves}
    print("阀门影响评估:")
    for valve, impact in impacts.items():
        print(f"  阀门 {valve}: 影响值={impact}")
  
    # 步骤3：构建优化模型
    prob = LpProblem("Valve_Optimization", LpMinimize)
  
    # 定义决策变量
    valve_vars = {valve: LpVariable(f"valve_{valve}", 0, 1, cat='Binary') 
                 for valve in critical_valves}
  
    # 目标函数：最小化总影响
    prob += lpSum(impacts[valve] * (1 - valve_vars[valve]) 
                 for valve in critical_valves)
  
    # 约束：所有泄漏点必须被隔离
    for i, leak in enumerate(leak_nodes):
        leak_paths = get_leak_paths(G, leak)
        for j, path in enumerate(leak_paths):
            # 每条路径至少关闭一个阀门
            prob += lpSum(1 - valve_vars[v] for v in path) >= 1, \
                   f"Leak{leak}_Path{j}_Constraint"
  
    # 步骤4：求解优化问题
    prob.solve(PULP_CBC_CMD(msg=False))
  
    # 解析结果
    shutdown_valves = [v for v in critical_valves if valve_vars[v].value() == 0]
    print(f"\n优化结果状态: {LpStatus[prob.status]}")
    print(f"需关闭的阀门: {shutdown_valves}")
    print(f"总影响值: {value(prob.objective)}")
  
    return shutdown_valves

# 执行示例
if __name__ == "__main__":
    # 构建管网并设置泄漏点
    G = build_gas_network()
    leak_nodes = ["L1"]
  
    # 运行优化
    shutdown_valves = optimize_valve_shutdown(G, leak_nodes)